import numpy as np
import json
from collections import defaultdict
import os, shutil

data_path = "/Users/pitaloveu/working_data/re_id_seq" 
retrieve_file = "tc_result.txt"

dict_tree = lambda : defaultdict( dict_tree )

def get_person_ID( dir_name ):
    """
    dir_name is in a form: 
        personID-UnknownFeatures
    """
    return dir_name.split( '-' )[0]

def generate_json_from_directory( d ):
    """
    input the directory contains all the .jpgs
    in a form:

    /d
        id1
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
        id2
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            ....
            frameID_x_y_h_w_cameraID.jpg
        ...
        id120
            frameID_x_y_h_w_cameraID.jpg
            frameID_x_y_h_w_cameraID.jpg
            ....
            frameID_x_y_h_w_cameraID.jpg
    """
    # first fetch all the directories in d

    all_files = os.listdir( d )
    all_dirs  = [ v for v in all_files if os.path.isdir( os.path.join( data_path, v ) ) ]

    personID_list = list( map( get_person_ID , all_dirs ) )

    result_json = dict_tree()

    for personDir in all_dirs:
        personID = get_person_ID( personDir )

        all_jpgs = os.listdir( os.path.join( data_path , personDir ) )
        all_jpgs = [ j for j in all_jpgs \
                if os.path.isfile( os.path.join(data_path , personDir, j) ) ]

        for jpgName in all_jpgs:
            jpgNameList = jpgName.split( '_' )
            frameID = jpgNameList[0]
            x = jpgNameList[1]
            y = jpgNameList[2]
            h = jpgNameList[3]
            w = jpgNameList[4]
            cameraID = jpgNameList[5]
            ## recently we use random number
            cameraID = str (np.random.rand(1)[0])

            result_json[personID][cameraID]["bbox_info"][frameID]["bottom"] = x
            result_json[personID][cameraID]["bbox_info"][frameID]["top"] = y
            result_json[personID][cameraID]["bbox_info"][frameID]["left"] = h
            result_json[personID][cameraID]["bbox_info"][frameID]["right"] = w

    data = json.dumps( result_json )
    print ( data )
    f = open( 'result.json' , 'w' )
    f.write(  data )
    f.close()

def generate_json_from_txt( txt ):
    """
    the txt file contains all the possible ID information in lines

    ID1 : jp1_name , jp2_name, jp3_name.........
    ID2 : jp1_name , jp2_name, jp3_name.........
    ....
    IDn : jp1_name , jp2_name, jp3_name.........

    here is the solution, first read every line from the txt into a list
    then , split by ":" to get the id and the jpgS

    then split the jpgS by ","

    then split the jpg by "_" to get all the information
    """

    fo = open( txt , "r" )
    source = fo.readlines()
    result_json = dict_tree()

    for line in source:
        personID = line.split( ':' )[0]
        jpg_info = line.split( ':' )[1]
        jpgs = jpg_info.split( ',' )

        for jpg in jpgs:
            infos = jpg.split( '/' )[-1]
            json_infos = infos.split( '_' )
            frameID = json_infos[0]
            x = json_infos[1]
            y = json_infos[2]
            w = json_infos[3]
            h = json_infos[4]
            cameraID = json_infos[5].replace( '*' , '/' )
            width = json_infos[6]
            heigh = json_infos[7]

            left = float( x ) / float( width )
            right = left + float( w ) / float( width )
            top = 1. - ( float( y ) / float( heigh ) )
            bottom = top - float( h ) / float( heigh )

            result_json[personID][cameraID]['bbox_info'][frameID]['bottom'] = bottom
            result_json[personID][cameraID]['bbox_info'][frameID]['top'] = top
            result_json[personID][cameraID]['bbox_info'][frameID]['left'] = left
            result_json[personID][cameraID]['bbox_info'][frameID]['right'] = right

    data = json.dumps( result_json )
    f = open( 'result.json' , 'w' )
    f.write(  data )
    f.close()

#generate_json_from_directory( data_path )
generate_json_from_txt( retrieve_file )
